from threading import Lock
import os

class ChunkFile(object):
    def __init__(self, path, blocksize, size=None):
        super(ChunkFile, self).__init__()
        self._path = path
        self._blocksize = blocksize

        # Fetch the file size if none is given.
        if size == None:
            # Notice that this will throw an os.error if the file does not exist.
            self._size = os.path.getsize(path)
        else:
            self._size = size

        # Calculate the size of the last block.
        self._last_blocksize = blocksize if self._size % blocksize == 0 else self._size % blocksize

        # ... and the total number of blocks.
        self._chunks = self._size / blocksize + (0 if self._last_blocksize == blocksize else 1)

        # We maintain a lock to be sure that the file position is trustworthy.
        self._lock = Lock()

        # Open the file (create it if it does not exist).
        if not os.path.exists(path):
            self._file = open(self._path, 'w+b')
        else:
            self._file = open(self._path, 'r+b')

    def write_chunk(self, idx, data):
        # Start by validating the inputs.
        if idx > self._chunks-1 or idx < 0:
            raise Exception('Invalid index given: %i'%idx)
        if idx == self._chunks-1:
            if len(data) != self._last_blocksize:
                raise Exception('Data of invalid length given for the last chunk: %i'%len(data))
        else:
            if len(data) != self._blocksize:
                raise Exception('Data of invalid size given: %i'%len(data))
        
        # Write the given data at the specified index.
        with self._lock:
            self._file.seek(self._blocksize * idx)
            self._file.write(data)

    def read_chunk(self, idx):
        # Start by validating the inputs.
        if idx > self._chunks-1 or idx < 0:
            raise Exception('Invalid index given: %i'%idx)

        # Find out how many bytes we should read.
        need_to_read = self._blocksize if idx != self._chunks - 1 else self._last_blocksize
        
        # Read the chunk data into memory.
        pieces = []
        read_so_far = 0
        with self._lock:
            self._file.seek(self._blocksize * idx)
            while read_so_far < need_to_read:
                piece = self._file.read(need_to_read - read_so_far)
                pieces.append(piece)
                read_so_far += len(piece)
        return ''.join(pieces)

    def close(self):
        with self._lock:
            self._file.close()

if __name__ == '__main__':
    test = ChunkFile('hest.test', 16, 1000)
    for x in range(0, 62):
        test.write_chunk(x, '1234567890123456')
    test.write_chunk(62, 'xxxxxxxx')
    test.write_chunk(10, 'yyyyyyyyyyyyyyyy')
    
    if test.read_chunk(62) != 'xxxxxxxx':
        print 'Error reading the last chunk.'
    if test.read_chunk(10) != 'yyyyyyyyyyyyyyyy':
        print 'Error reading chunk 10.'
    if test.read_chunk(0) != '1234567890123456':
        print 'Error reading the first chunk.'

    test.close()